package ai.zixing.mashibing.basic_class.class10;

import java.util.Collection;
import java.util.Comparator;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.PriorityQueue;
import java.util.Set;
import java.util.Stack;

/**
 * 最小生成树算法之Kruskal
 */
public class Code04_Kruskal {

    public static class UnionFind {

        private HashMap<Node, Node> fatherMap;

        private HashMap<Node, Integer> sizeMap;

        public UnionFind() {
            this.fatherMap = new HashMap<>();
            this.sizeMap = new HashMap<>();
        }

        public void makeSets(Collection<Node> nodes) {
            fatherMap.clear();
            sizeMap.clear();
            for (Node node : nodes) {
                fatherMap.put(node, node);
                sizeMap.put(node, 1);
            }
        }

        public Node findFather(Node node) {
            Stack<Node> path = new Stack<>();
            while (node != fatherMap.get(node)) {
                path.add(node);
                node = fatherMap.get(node);
            }
            while (!path.isEmpty()) {
                fatherMap.put(path.pop(), node);
            }
            return node;
        }

        public boolean isSameSet(Node node1, Node node2) {
            return findFather(node1) == findFather(node2);
        }

        public void union(Node a, Node b) {
            if (a == null || b == null) {
                return;
            }
            Node aHead = findFather(a);
            Node bHead = findFather(b);
            if (aHead != bHead) {
                int aSetSize = sizeMap.get(aHead);
                int bSetSize = sizeMap.get(bHead);
                Node bigNode = aSetSize >= bSetSize ? aHead : bHead;
                Node smallNode = aSetSize < bSetSize ? aHead : bHead;
                fatherMap.put(smallNode, bigNode);
                sizeMap.put(bigNode, aSetSize + bSetSize);
                sizeMap.remove(smallNode);
            }
        }
    }

    public static class EdgeComparator implements Comparator<Edge> {
        @Override
        public int compare(Edge o1, Edge o2) {
            return o1.weight - o2.weight;
        }
    }

    public static Set<Edge> kruskalMST(Graph graph) {
        // 初始化图
        UnionFind unionFind = new UnionFind();
        unionFind.makeSets(graph.nodes.values());
        // 权重排序
        PriorityQueue<Edge> priorityQueue = new PriorityQueue<>(new EdgeComparator());
        // 将图的所有边存入堆 M条边 O(logM)
        priorityQueue.addAll(graph.edges);
        Set<Edge> result = new HashSet<>();
        // M条边 O(logM)
        while (!priorityQueue.isEmpty()) {
            Edge edge = priorityQueue.poll();
            // O(1)
            if (!unionFind.isSameSet(edge.from, edge.to)) {
                result.add(edge);
                unionFind.union(edge.from, edge.to);
            }
        }
        return result;
    }

}
